package zunge.lbe.analysis;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;

import zunge.tools.Tools;

public class DataSetFold {
	public DataSet dataSet;
	public int v;	
	
	public String folder;
	public static final double Ratio = 0.2; // For each test user, 20% of the apps are known.
	
	public String getFile(String label){return folder +label;}
	public static final String fsuf_train = "train.txt";
	public static final String fsuf_test = "test.txt";
	
	
	public static List<DataSetFold> getFolds(DataSet dataSet, int V, boolean create) throws Exception{
		
		List<Integer> userList = new ArrayList<Integer>(dataSet.uMap.values());
		Collections.shuffle(userList);
		List<HashSet<Integer>> foldTestUsers = new ArrayList<HashSet<Integer>>();
		for(int v = 0;v < V;++v){foldTestUsers.add(new HashSet<Integer>());}
		if(create){
			for(int i = 0;i < userList.size();++i){
				foldTestUsers.get(i%V).add(userList.get(i));
			}
		}
		
		List<DataSetFold> folds = new ArrayList<DataSetFold>();
		for(int v = 0;v < V;++v){folds.add(new DataSetFold(dataSet, v, foldTestUsers.get(v), create));}
		
		return folds;
	}
	
	
	DataSetFold(DataSet dataSet, int v, HashSet<Integer> testUserSet, boolean create) throws Exception{
		this.dataSet = dataSet;
		this.v = v;
		this.folder = dataSet.folder + "fold_"+v+"/";		
		
		if(create){
			new File(folder).mkdir();
			
			//Generate Train and Test file
			BufferedReader br = new BufferedReader(new FileReader(dataSet.getFile(dataSet.fsuf_recordU)));
			BufferedWriter bwtrain = new BufferedWriter(new FileWriter(getFile(fsuf_train)));
			BufferedWriter bwtest = new BufferedWriter(new FileWriter(getFile(fsuf_test)));
			
			String line; int uid;
			while((line=br.readLine())!=null){
				String[] frags = line.split(" ");
				uid = Integer.parseInt(frags[0]);
				
				String[] apds = frags[1].split(",");
				
				HashSet<Integer> aSet = new HashSet<Integer>();
				if(apds.length==1){continue;}
				for(String apd:apds){
					int aid = Integer.parseInt(apd.substring(0, apd.indexOf(":")));
					aSet.add(aid);
				}
				List<Integer> aList = new ArrayList<Integer>(aSet);
				HashSet<Integer> trainAppIndices = Tools.getRandomSelections(aList.size(), Ratio);
				HashSet<Integer> trainASet = new HashSet<Integer>();
				for(int i = 0;i < aList.size();++i){if(trainAppIndices.contains(i)){trainASet.add(aList.get(i));}}
				
				
				if(testUserSet.contains(uid)){ //Test Set
					for(String apd:apds){
						String[] apdFrags = apd.split(":");
						int aid = Integer.parseInt(apdFrags[0]);
						if(trainASet.contains(aid)){//Train
							bwtrain.write(frags[0]+" "+apdFrags[0]+" "+apdFrags[1]+" "+apdFrags[2]+"\n");
						}else{
							bwtest.write(frags[0]+" "+apdFrags[0]+" "+apdFrags[1]+" "+apdFrags[2]+"\n");
						}
					}
				}
				else{ //Train Set
					for(String apd:apds){
						String[] apdFrags = apd.split(":");
						bwtrain.write(frags[0]+" "+apdFrags[0]+" "+apdFrags[1]+" "+apdFrags[2]+"\n");
					}
				}
			}
			
			br.close();
			bwtrain.close();
			bwtest.close();
		}
		
	}
}
